package com.rrd.hive.udf.brickhouse.udf.collect;

import java.util.LinkedHashMap;
import java.util.Map;
import java.util.TreeSet;

import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.parse.SemanticException;
import org.apache.hadoop.hive.ql.udf.generic.AbstractGenericUDAFResolver;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFEvaluator;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.StandardMapObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.WritableConstantIntObjectInspector;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
import org.apache.log4j.Logger;

@Description(name = "collect_max",
        value = "_FUNC_(x, val, n) - Returns an map of the max N numeric values in the aggregation group ")
public class CollectMaxUDAF extends AbstractGenericUDAFResolver {
    public static final Logger LOG = Logger.getLogger(CollectMaxUDAF.class);
    public static int DEFAULT_MAX_VALUES = 20;

    @Override
    public GenericUDAFEvaluator getEvaluator(TypeInfo[] parameters)
            throws SemanticException {
        return new MapCollectMaxUDAFEvaluator();
    }

    public static class MapCollectMaxUDAFEvaluator
            extends GenericUDAFEvaluator {
        // For PARTIAL1 and COMPLETE: ObjectInspectors for original data
        private PrimitiveObjectInspector inputKeyOI;
        private PrimitiveObjectInspector inputValOI;
        // For PARTIAL2 and FINAL: ObjectInspectors for partial aggregations (list
        // of objs)
        private StandardMapObjectInspector internalMergeOI;
        private boolean descending = true;
        private int numValues = CollectMaxUDAF.DEFAULT_MAX_VALUES;

        public MapCollectMaxUDAFEvaluator() {
            this(true);
        }

        public MapCollectMaxUDAFEvaluator(boolean desc) {
            this.descending = desc;
        }

        public class SortedKeyValue implements Comparable {
            private Object keyObj;
            private Object valObj;

            public SortedKeyValue(Object keyObj, Object valObj) {
                this.keyObj = keyObj;
                this.valObj = valObj;
            }

            @Override
            public boolean equals(Object other) {
                if (!(other instanceof SortedKeyValue)) {
                    return false;
                }
                SortedKeyValue otherKV = (SortedKeyValue) other;
                if (this.getKey().equals(otherKV.getKey())) {
                    return true;
                } else {
                    return false;
                }
            }

            public Object getKey() {
                return this.keyObj == null ? null
                    : MapCollectMaxUDAFEvaluator.this.inputKeyOI
                        .getPrimitiveJavaObject(this.keyObj);
            }

            public Double getValue() {
                return this.valObj == null ? null
                    : ((Number) MapCollectMaxUDAFEvaluator.this.inputValOI
                        .getPrimitiveJavaObject(this.valObj)).doubleValue();
            }

            @Override
            public String toString() {
                return this.getKey() + ":=" + this.getValue();
            }

            @Override
            public int compareTo(Object arg1) {
                SortedKeyValue otherKV = (SortedKeyValue) arg1;
                int cmp = this.compareKV(otherKV);
                return MapCollectMaxUDAFEvaluator.this.descending ? cmp
                    : -1 * cmp;
            }

            public int compareKV(SortedKeyValue otherKV) {

                double thisNumber = this.getValue();
                double otherNumber = otherKV.getValue();
                int sigNum = (int) Math.signum(otherNumber - thisNumber);
                if (sigNum != 0) {
                    return sigNum;
                } else {
                    return ((Comparable) this.getKey())
                        .compareTo(otherKV.getKey());
                }
            }
        }

        class MapAggBuffer implements AggregationBuffer {
            private TreeSet<SortedKeyValue> sortedValues = new TreeSet<>();

            public void addValue(Object keyObj, Object valObj) {
                if (this.sortedValues
                    .size() < MapCollectMaxUDAFEvaluator.this.numValues) {
                    SortedKeyValue newValue = new SortedKeyValue(
                        MapCollectMaxUDAFEvaluator.this.inputKeyOI
                            .copyObject(keyObj),
                        MapCollectMaxUDAFEvaluator.this.inputValOI
                            .copyObject(valObj));
                    this.sortedValues.add(newValue);
                } else {
                    SortedKeyValue minValue = this.sortedValues.last();
                    SortedKeyValue biggerValue = new SortedKeyValue(
                        MapCollectMaxUDAFEvaluator.this.inputKeyOI
                            .copyObject(keyObj),
                        MapCollectMaxUDAFEvaluator.this.inputValOI
                            .copyObject(valObj));
                    int cmp = biggerValue.compareTo(minValue);
                    if (cmp < 0) {
                        this.sortedValues.remove(minValue);
                        this.sortedValues.add(biggerValue);
                    }
                }
            }

            public Map getValueMap() {
                LinkedHashMap<Object, Object> reverseOrderMap = new LinkedHashMap<>();
                for (SortedKeyValue kv : this.sortedValues) {
                    reverseOrderMap.put(kv.keyObj, kv.valObj);
                }
                return reverseOrderMap;
            }

            public void reset() {
                this.sortedValues.clear();
            }
        }

        @Override
        public ObjectInspector init(Mode m, ObjectInspector[] parameters)
                throws HiveException {
            super.init(m, parameters);
            CollectMaxUDAF.LOG
                .error(" CollectMaxUDAF.init() - Mode= " + m.name());
            for (int i = 0; i < parameters.length; ++i) {
                CollectMaxUDAF.LOG
                    .error(" ObjectInspector[ " + i + " ] = " + parameters[0]);
            }
            if (parameters.length > 2) {
                if (parameters[2] instanceof WritableConstantIntObjectInspector) {
                    WritableConstantIntObjectInspector nvOI = (WritableConstantIntObjectInspector) parameters[2];
                    this.numValues = nvOI.getWritableConstantValue().get();
                    CollectMaxUDAF.LOG
                        .info(" Setting number of values to " + this.numValues);
                } else {
                    throw new HiveException(
                        "Number of values must be a constant int.");
                }
            }

            // init output object inspectors
            // The output of a partial aggregation is a map
            if (m == Mode.PARTIAL1) {
                this.inputKeyOI = (PrimitiveObjectInspector) parameters[0];
                this.inputValOI = (PrimitiveObjectInspector) parameters[1];
                return ObjectInspectorFactory.getStandardMapObjectInspector(
                    this.inputKeyOI, this.inputValOI);
            } else {
                if (!(parameters[0] instanceof StandardMapObjectInspector)) {
                    this.inputKeyOI = (PrimitiveObjectInspector) parameters[0];
                    this.inputValOI = (PrimitiveObjectInspector) parameters[1];
                } else {
                    this.internalMergeOI = (StandardMapObjectInspector) parameters[0];
                    this.inputKeyOI = (PrimitiveObjectInspector) this.internalMergeOI
                        .getMapKeyObjectInspector();
                    this.inputValOI = (PrimitiveObjectInspector) this.internalMergeOI
                        .getMapValueObjectInspector();
                }
            }
            return ObjectInspectorFactory.getStandardMapObjectInspector(
                this.inputKeyOI, this.inputValOI);
        }

        @Override
        public AggregationBuffer getNewAggregationBuffer()
                throws HiveException {
            MapAggBuffer buff = new MapAggBuffer();
            this.reset(buff);
            return buff;
        }

        @Override
        public void iterate(AggregationBuffer agg, Object[] parameters)
                throws HiveException {
            Object k = parameters[0];
            Object v = parameters[1];
            if (k == null || v == null) {
                throw new HiveException(
                    "Key or value is null.  k = " + k + " , v = " + v);
            }

            if (k != null) {
                MapAggBuffer myagg = (MapAggBuffer) agg;

                this.putIntoSet(k, v, myagg);
            }
        }

        @Override
        public void merge(AggregationBuffer agg, Object partial)
                throws HiveException {
            MapAggBuffer myagg = (MapAggBuffer) agg;
            Map<Object, Object> partialResult = (Map<Object, Object>) this.internalMergeOI
                .getMap(partial);
            for (Object i : partialResult.keySet()) {
                this.putIntoSet(i, partialResult.get(i), myagg);
            }
        }

        @Override
        public void reset(AggregationBuffer buff) throws HiveException {
            MapAggBuffer arrayBuff = (MapAggBuffer) buff;
            arrayBuff.reset();
        }

        @Override
        public Object terminate(AggregationBuffer agg) throws HiveException {
            MapAggBuffer myagg = (MapAggBuffer) agg;
            return myagg.getValueMap();

        }

        private void putIntoSet(Object key, Object val, MapAggBuffer myagg) {
            myagg.addValue(key, val);
        }

        @Override
        public Object terminatePartial(AggregationBuffer agg)
                throws HiveException {

            MapAggBuffer myagg = (MapAggBuffer) agg;
            Map<Object, Object> vals = myagg.getValueMap();
            return vals;
        }
    }

}
